In [1]:
from __future__ import division, print_function
# Third-party
from astropy import log as logger
import emcee
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm
from scipy.misc import logsumexp
import triangle
%matplotlib inline
# Custom
import streams.coordinates as stc
import streamteam.io as io
import streamteam.dynamics as sd
import streamteam.potential as sp
from streamteam.units import galactic
from streams.rewinder.likelihood import rewinder_likelihood
In [43]:
scf = io.APWSCFReader("/Users/adrian/projects/morphology/simulations/runs/spherical/")
ppars = scf.read_potential(units=galactic)
potential = sp.LeeSutoNFWPotential(v_h=ppars['v_h'], r_h=ppars['r_h'],
a=1., b=1., c=1., units=galactic)
# read all particles
tbl = scf.read_snap("SNAP006", units=galactic)
ptcl = io.tbl_to_w(tbl)
unbound_tbl = tbl[tbl['tub']!=0]
unbound_stars = ptcl[tbl['tub']!=0]
prog = np.median(ptcl[tbl['tub']==0], axis=0).reshape(1,6)
print("{} unbound stars".format(len(unbound_stars)))
np.random.seed(42)
stars_ix = np.random.randint(len(unbound_stars), size=512)
stars_ix = np.ones(len(unbound_stars)).astype(bool) # HACK
stars_tbl = tbl[tbl['tub']!=0][stars_ix]
stars = unbound_stars[stars_ix]
print(stars.shape, potential.parameters)
sat_mass = 2.5E6
In [39]:
# progenitor orbit
t,prog_w = potential.integrate_orbit(prog.copy(), dt=-1., nsteps=6000)
# orbits of stars that come unbound
ub_t,ub_w = potential.integrate_orbit(unbound_stars.copy(), dt=-1., nsteps=6000)
In [40]:
betas = np.ones(len(stars_ix))
betas[stars[:,1] > (stars[:,0] + 21.6)] = -1.
fig = sd.plot_orbits(prog_w, marker=None, subplots_kwargs=dict(sharex=True, sharey=True))
fig = sd.plot_orbits(unbound_stars, linestyle='none', alpha=0.4, axes=fig.axes)
fig = sd.plot_orbits(prog, linestyle='none', marker='+', markeredgewidth=2.,
markersize=15, axes=fig.axes)
fig = sd.plot_orbits(stars[betas>0], linestyle='none', ms=6,
marker='o', color='g', axes=fig.axes, alpha=0.5)
fig = sd.plot_orbits(stars[betas<0], linestyle='none', ms=6,
marker='o', color='r', axes=fig.axes, alpha=0.5, )
x = np.linspace(-50,50,100)
fig.axes[0].plot(x, x+21.6, marker=None)
Out[40]:
In [44]:
ll = rewinder_likelihood(-1., 6000.,
potential.c_instance,
prog, stars[:16],
2.5E6, 0.,
1., betas[:16], -0.3)
print(ll)
In [6]:
fig = sd.plot_orbits(ub_w[:,:100], marker=None, alpha=0.1)
In [7]:
def construct_basis(prog_w, theta=0.):
x1_hat = prog_w[...,:3].copy()
x1_hat /= np.linalg.norm(x1_hat, axis=-1)[...,None]
x3_hat = np.cross(prog_w[...,:3], prog_w[...,3:])
x3_hat /= np.linalg.norm(x3_hat, axis=-1)[...,None]
x2_hat = np.cross(x3_hat, x1_hat)
x2_hat /= np.linalg.norm(x2_hat, axis=-1)[...,None]
tmp = x1_hat[...,0]
x1_hat[...,0] = x1_hat[...,0]*np.cos(theta) + x1_hat[...,1]*np.sin(theta)
x1_hat[...,1] = -tmp*np.sin(theta) + x1_hat[...,1]*np.cos(theta)
tmp = x2_hat[...,0]
x2_hat[...,0] = x2_hat[...,0]*np.cos(theta) + x2_hat[...,1]*np.sin(theta)
x2_hat[...,1] = -tmp*np.sin(theta) + x2_hat[...,1]*np.cos(theta)
return x1_hat, x2_hat, x3_hat
def change_basis(w, prog_w, theta=0.):
x1_hat,x2_hat,x3_hat = construct_basis(prog_w, theta)
# translate to progenitor position
dx = w[...,:3] - prog_w[...,:3]
dv = w[...,3:] - prog_w[...,3:]
# rotate into basis
x1 = dx[...,0]*x1_hat[...,0] + dx[...,1]*x1_hat[...,1] + dx[...,2]*x1_hat[...,2]
x2 = dx[...,0]*x2_hat[...,0] + dx[...,1]*x2_hat[...,1] + dx[...,2]*x2_hat[...,2]
x3 = dx[...,0]*x3_hat[...,0] + dx[...,1]*x3_hat[...,1] + dx[...,2]*x3_hat[...,2]
x = np.array([x1.T,x2.T,x3.T]).T
return x
In [8]:
rtides = np.sqrt(np.sum(prog_w[:,0,:3]**2,axis=-1))*(2.5E6 / potential.c_instance.mass_enclosed(prog_w[:,0,:3].copy()))**(1/3.)
In [9]:
plt.plot(potential.c_instance.mass_enclosed(prog_w[:,0,:3].copy()))
Out[9]:
In [10]:
RR = np.sqrt(np.sum(prog_w[:,0,:3]**2,axis=-1))
In [11]:
plt.plot(RR, rtides, marker=None)
Out[11]:
In [45]:
xprime = change_basis(ub_w, prog_w, theta=-0.3)
ixes = (6000 - unbound_tbl['tub']).astype(int)
at_tub = np.array([xprime[ix,i] for i,ix in enumerate(ixes)])
rtides_tub = np.array([rtides[ix] for i,ix in enumerate(ixes)])
ltz = at_tub[:,0] < 0.
at_tub[ltz,0] += 1.2*rtides[ltz]
at_tub[~ltz,0] -= 1.2*rtides[~ltz]
fig = sd.plot_orbits(at_tub, marker='.', linestyle='none')
# fig = sd.plot_orbits(at_tub[stars_ix], marker='o', color='r', linestyle='none', axes=fig.axes)
for ax in fig.axes:
ax.set_xlim(-3,3)
ax.set_ylim(-3,3)
In [46]:
vhs = np.linspace(0.3,0.7,25)
l = []
for v_h in vhs:
p = sp.LeeSutoNFWPotential(v_h=v_h, r_h=20, a=1, b=1, c=1, units=galactic)
ll = rewinder_likelihood(-1., 6000,
p.c_instance,
prog, stars[:8],
2.5E6, 0.,
1.2, betas, -0.3)
b = np.ones((6000,1))
b[1] = b[-1] = 0.5
l.append(logsumexp(ll, axis=0, b=b).sum())
l = np.array(l)
fig,axes = plt.subplots(1,2,figsize=(10,5))
axes[0].plot(vhs, l, marker='o')
axes[1].plot(vhs, np.exp(l - l.max()), marker='o')
Out[46]:
Why is the log-likelihood -inf for so many values? Have to turn on debugging to get orbits back, then compare orbits...
In [14]:
# p = sp.LeeSutoNFWPotential(v_h=0.4, r_h=20, a=1, b=1, c=1, units=galactic)
# ll,x,v = rewinder_likelihood(6000., 0., -1.,
# p.c_instance,
# prog, stars[:4],
# 2.5E6, 0.,
# 1., betas, -0.3)
# print(ll)
# fig = sd.plot_orbits(x, ix=0, marker=None)
# fig = sd.plot_orbits(x, ix=3, marker=None, axes=fig.axes)
# fig.axes[1].set_ylim(-10,10)
# fig.axes[2].set_ylim(-10,10)
Aha, it's because the orbits really do look different at lower v_h...
Now I'll observe some stars to have some fake observed data to play with
In [118]:
stars_hel = stc.gal_to_hel(stars)
sigma = np.zeros_like(stars_hel)
sigma[:,:2] = 1E-8
sigma[:,2] = stars_hel[:,2] * 0.02
sigma[:,3:5] = ((10.*u.km/u.s) / (30*u.kpc)).to(u.radian/u.Myr, equivalencies=u.dimensionless_angles()).value
sigma[:,5] = (10*u.km/u.s).to(u.kpc/u.Myr).value
sigma = np.abs(sigma)
obs_stars_hel = np.random.normal(stars_hel, sigma)
obs_stars = stc.hel_to_gal(obs_stars_hel)
# -----------------------------------------------------
prog_hel = stc.gal_to_hel(prog)
sigma_prog = np.zeros_like(prog_hel)
sigma_prog[:,:2] = 1E-8
sigma_prog[:,2] = prog_hel[:,2] * 0.02
# sigma_prog[:,3:5] = ((10.*u.km/u.s) / (30*u.kpc)).to(u.radian/u.Myr, equivalencies=u.dimensionless_angles()).value # GAIA
sigma_prog[:,3:5] = ((1.*u.km/u.s) / (30*u.kpc)).to(u.radian/u.Myr, equivalencies=u.dimensionless_angles()).value # TESTING
# sigma_prog[:,5] = (10*u.km/u.s).to(u.kpc/u.Myr).value # RR Lyr?
sigma_prog[:,5] = (1*u.km/u.s).to(u.kpc/u.Myr).value
sigma_prog = np.abs(sigma_prog)
obs_prog_hel = np.random.normal(prog_hel, sigma_prog)
obs_prog = stc.hel_to_gal(obs_prog_hel)
Save data to test file
In [48]:
# names = 'x y z vx vy vz l b d mul mub vr tub tail'
# np.savetxt("../streams/rewinder/tests/true_stars.txt",
# np.hstack((stars, stars_hel, stars_tbl['tub'][:,np.newaxis], betas[:,np.newaxis])),
# header=names)
# names = 'x y z vx vy vz l b d mul mub vr err_l err_b err_d err_mul err_mub err_vr tub tail'
# np.savetxt("../streams/rewinder/tests/obs_stars.txt",
# np.hstack((obs_stars, obs_stars_hel, sigma, stars_tbl['tub'][:,np.newaxis], betas[:,np.newaxis])),
# header=names)
# names = 'x y z vx vy vz l b d mul mub vr mass'
# np.savetxt("../streams/rewinder/tests/true_prog.txt",
# np.hstack((prog, prog_hel, np.array([[2.5E6]]))),
# header=names)
# names = 'x y z vx vy vz l b d mul mub vr err_l err_b err_d err_mul err_mub err_vr mass'
# np.savetxt("../streams/rewinder/tests/obs_prog.txt",
# np.hstack((obs_prog, obs_prog_hel, sigma_prog, np.array([[2.5E6]]))),
# header=names)
For each star, I need to either:
1)
or 2)
In [49]:
plt.figure(figsize=(8,6))
bins = np.linspace(0,0.02,50)
plt.hist(np.abs(potential.energy(stars[:,:3], stars[:,3:]) - potential.energy(prog[:,:3], prog[:,3:])),
bins=bins, label="all stars", normed=True, alpha=0.4);
pos = stc.hel_to_gal(np.random.normal(obs_stars_hel[0], sigma[0], size=(1000,6)))
plt.hist(np.abs(potential.energy(pos[:,:3], pos[:,3:]) - potential.energy(prog[:,:3], prog[:,3:])),
bins=bins, label="samples from one error ellipse", normed=True, alpha=0.4);
plt.xlabel("$\Delta E$")
plt.legend()
Out[49]:
In [234]:
# Energy
prog_E = potential.energy(prog[:,:3], prog[:,3:])
stars_dE = potential.energy(stars[:,:3], stars[:,3:]) - prog_E
mean_dE = np.mean(np.abs(stars_dE))
std_dE = np.std(np.abs(stars_dE))
print(mean_dE, std_dE)
print("Prog. E:", prog_E)
# Angular momentum -- can compute at t=0 because spherical
prog_J = np.cross(prog[:,:3], prog[:,3:])
prog_Jmag = np.sqrt(np.sum(prog_J**2, axis=-1))
stars_dJ = np.cross(stars[:,:3], stars[:,3:]) - prog_J
mean_dJ = np.mean(np.abs(stars_dJ), axis=0)
std_dJ = np.std(np.abs(stars_dJ), axis=0)
print(mean_dJ, std_dJ)
print("Prog. J:", prog_J)
I estimate the spread in energies...
In [225]:
Menc = potential.mass_enclosed(prog_w[:,0,:3])
E_scale = (sat_mass/Menc)**(1./3.) * np.sum(prog_w[:,0,3:]**2, axis=-1)
# E_scale = (sat_mass/Menc)**(1./3.) * np.sum(prog[:,3:]**2, axis=-1)
J_scale = (sat_mass/Menc)**(1./3.) * prog_J
_dJ = np.mean(J_scale)
_dE = np.mean(E_scale) # , np.std(E_scale)
_dE
Out[225]:
In [215]:
fig,axes = plt.subplots(1,2,figsize=(12,5))
a = 1.3
b = 2*a
bins = np.linspace(0,0.005,50)
axes[0].hist(np.abs(stars_dE), bins=bins, normed=True, alpha=0.4);
axes[0].axvline(a*_dE)
axes[0].axvline(b*_dE)
axes[1].hist(np.abs(stars_dJ), bins=50, normed=True, alpha=0.4);
axes[1].axvline(a*_dJ)
axes[1].axvline(b*_dJ)
Out[215]:
In [237]:
np.random.seed(42)
# First, try for a single star!
i = 2
# --------------------------------------------------------------
tot_nsamples = 100000
importance_samples_hel = np.zeros((tot_nsamples,6))
importance_samples_hel[...,:2] = obs_stars_hel[i,:2]
importance_samples_hel[...,2:] = np.random.normal(obs_stars_hel[i,2:], sigma[i,2:], size=(tot_nsamples,4))
importance_samples = stc.hel_to_gal(importance_samples_hel)
star_dE = np.abs(potential.energy(importance_samples[:,:3], importance_samples[:,3:]) - prog_E)
star_J = np.cross(importance_samples[:,:3], importance_samples[:,3:])
star_dJ = np.abs(np.sqrt(np.sum(star_J**2, axis=-1)) - prog_Jmag)
Eix = np.where((star_dE > (a*_dE)) & (star_dE < (b*_dE)))[0]# &
# (star_dJ > (a*_dJ)) & (star_dJ < (b*_dJ)))[0]
In [226]:
plt.plot(importance_samples_hel[:,4], importance_samples_hel[:,5], alpha=0.4, linestyle='none')
plt.plot(good_samples_hel[:,4], good_samples_hel[:,5], alpha=0.4, linestyle='none')
# plt.hist(good_samples_hel[:,4], normed=True, alpha=0.4)
Out[226]:
In [206]:
nsamples = 2048
good_samples = importance_samples[Eix][:nsamples]
good_samples_hel = stc.gal_to_hel(good_samples)
if len(good_samples) != nsamples:
raise ValueError()
# --------------------------------------------------------------
vhs = np.linspace(0.4,0.6,25)
l = np.zeros_like(vhs)
for j,v_h in enumerate(vhs):
p = sp.LeeSutoNFWPotential(v_h=v_h, r_h=20, a=1, b=1, c=1, units=galactic)
ll = rewinder_likelihood(-1., 6000,
p.c_instance,
prog, good_samples,
2.5E6, 0.,
1.25, np.zeros(nsamples) + betas[i], -0.3)
# l[j] += logsumexp(ll - good_samples_logp) - np.log(len(good_samples))
l[j] += logsumexp(ll) - np.log(nsamples)
In [205]:
fig,axes = plt.subplots(1,2,figsize=(10,5))
fig.suptitle("N: {}".format(nsamples))
axes[0].plot(vhs, l, marker='o')
axes[1].plot(vhs, np.exp(l - l[np.isfinite(l)].max()), marker='o')
axes[1].axvline(0.5, alpha=0.2, color='r')
Out[205]:
In [ ]:
In [ ]:
In [ ]:
In [119]:
nstars = len(sigma)
nsamples = 10000
importance_samples_hel = np.zeros((nsamples,nstars,6))
importance_samples_hel[...,:2] = obs_stars_hel[:,:2]
importance_samples_hel[...,2:] = np.random.normal(obs_stars_hel[:,2:], sigma[:,2:], size=(nsamples,nstars,4))
importance_samples = stc.hel_to_gal(importance_samples_hel.reshape(nsamples*nstars,6)).reshape(nsamples,nstars,6)
In [121]:
i = 2
plt.hist(potential.energy(importance_samples[:,i,:3], importance_samples[:,i,3:]) - \
potential.energy(prog[:,:3], prog[:,3:]))
# 3 sigma
plt.axvline(-dE)
plt.axvline(-3*dE)
Out[121]:
In [122]:
min_nsamples = 256
In [103]:
i = 1
star_dE = potential.energy(importance_samples[:,i,:3], importance_samples[:,i,3:]) - \
potential.energy(prog[:,:3], prog[:,3:]) - betas[i]*dE*12/4
Eix = np.where((star_dE > (-3*dE*3/4)) & (star_dE < (3*dE*3/4)))[0][:min_nsamples]
if len(Eix) < min_nsamples:
raise ValueError()
good_samples = importance_samples[Eix,i]
good_samples_hel = stc.gal_to_hel(good_samples)
good_samples_logp = norm.logpdf(good_samples_hel[:,2:], loc=obs_stars_hel[i,2:], scale=sigma[i,2:]).sum(axis=-1)
In [287]:
t,sample_w = potential.integrate_orbit(good_samples, dt=-1., nsteps=6000)
In [288]:
fig = sd.plot_orbits(prog_w, marker=None)
fig = sd.plot_orbits(sample_w, marker=None, axes=fig.axes, alpha=0.1)
In [ ]:
In [ ]:
In [127]:
min_nsamples = 128
vhs = np.linspace(0.4,0.6,15)
l = np.zeros_like(vhs)
# for i in np.append(np.where(betas < 0.)[0][:4], np.where(betas > 0.)[0][:4]):
for i in np.where(betas > 0.)[0][:4]:
star_dE = np.abs(potential.energy(importance_samples[:,i,:3], importance_samples[:,i,3:]) - \
potential.energy(prog[:,:3], prog[:,3:]))
Eix = np.where((star_dE > dE) & (star_dE < 3*dE))[0]
print(len(Eix))
Eix = Eix[:min_nsamples]
good_samples = importance_samples[Eix,i]
good_samples_hel = stc.gal_to_hel(good_samples)
good_samples_logp = norm.logpdf(good_samples_hel[:,2:], loc=obs_stars_hel[i,2:], scale=sigma[i,2:]).sum(axis=-1)
for j,v_h in enumerate(vhs):
p = sp.LeeSutoNFWPotential(v_h=v_h, r_h=20, a=1, b=1, c=1, units=galactic)
ll = rewinder_likelihood(6000., 0., -1.,
p.c_instance,
prog, good_samples,
2.5E6, 0.,
2., np.zeros(min_nsamples) + betas[i], -0.3)
l[j] += logsumexp(ll - good_samples_logp) - np.log(len(good_samples))
#l[j] += logsumexp(ll) - np.log(len(good_samples))
In [128]:
fig,axes = plt.subplots(1,2,figsize=(10,5))
axes[0].plot(vhs, l, marker='o')
axes[1].plot(vhs, np.exp(l - l[np.isfinite(l)].max()), marker='o')
axes[1].axvline(0.5, alpha=0.2, color='r')
Out[128]:
In [ ]: